import os
import cv2
import imgaug as ia
from imgaug import augmenters as iaa
from lxml import etree
import shutil
import numpy as np
import re

# 定义输入输出路径
base_input_dir = r'D:\lingsih\custom_training~1\home\hdkj\test\aa\custom_training'  # 基础输入目录
input_image_dir = os.path.join(base_input_dir, 'JPEGImages')  # 输入图片目录
input_xml_dir = os.path.join(base_input_dir, 'Annotations')  # 输入XML目录

# 在基础输入目录下创建 newdata 目录
newdata_dir = os.path.join(base_input_dir, 'newdata')
os.makedirs(newdata_dir, exist_ok=True)

# 创建 newdata 下的 Annotations 和 JPEGImages 目录
output_image_dir = os.path.join(newdata_dir, 'JPEGImages')  # 输出图片目录
output_xml_dir = os.path.join(newdata_dir, 'Annotations')  # 输出XML目录
os.makedirs(output_image_dir, exist_ok=True)
os.makedirs(output_xml_dir, exist_ok=True)

# 数据增强序列
seq = iaa.Sequential([
    iaa.Fliplr(0.5),  # 水平翻转
    iaa.Affine(
        rotate=(-10, 10),  # 随机旋转
        scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},  # 随机缩放
        translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)}  # 随机平移
    ),
    iaa.Multiply((0.8, 1.2)),  # 改变亮度
    iaa.AddToHueAndSaturation((-20, 20)),  # 改变色调和饱和度
])

def parse_xml(xml_file):
    """
    解析XML文件并返回根元素和边界框信息。
    """
    tree = etree.parse(xml_file)
    root = tree.getroot()
    objects = []
    for obj in root.findall('object'):
        bbox = obj.find('bndbox')
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)
        objects.append({'name': obj.find('name').text, 'bbox': [xmin, ymin, xmax, ymax]})
    return root, objects


def update_xml(root, objects, bbs_aug):
    """
    更新XML文件中的边界框信息。
    """
    for obj, bb_aug in zip(objects, bbs_aug.bounding_boxes):
        bbox = obj['bbox']
        new_bbox = [int(bb_aug.x1), int(bb_aug.y1), int(bb_aug.x2), int(bb_aug.y2)]
        obj_elem = root.find(".//object[name='{}']".format(obj['name']))
        bndbox = obj_elem.find('bndbox')
        bndbox.find('xmin').text = str(new_bbox[0])
        bndbox.find('ymin').text = str(new_bbox[1])
        bndbox.find('xmax').text = str(new_bbox[2])
        bndbox.find('ymax').text = str(new_bbox[3])
    return root


def save_xml(root, output_xml_path):
    """
    将更新后的XML内容保存到指定路径。
    """
    with open(output_xml_path, 'wb') as f:
        f.write(etree.tostring(root, pretty_print=True))


def copy_original_files(input_image_dir, input_xml_dir, output_image_dir, output_xml_dir):
    """
    将原始文件复制到新的输出目录。
    """
    for file_name in os.listdir(input_image_dir):
        full_file_name = os.path.join(input_image_dir, file_name)
        if os.path.isfile(full_file_name):
            shutil.copy(full_file_name, output_image_dir)
    for file_name in os.listdir(input_xml_dir):
        full_file_name = os.path.join(input_xml_dir, file_name)
        if os.path.isfile(full_file_name):
            shutil.copy(full_file_name, output_xml_dir)


def get_max_index(image_files):
    """
    获取现有图片文件中的最大索引值。
    """
    indices = [int(re.search(r"(\d+)\.[^.]+$", f).group(1)) for f in image_files if re.search(r"(\d+)\.[^.]+$", f)]
    return max(indices) if indices else 0


def augment_images_and_xmls(input_image_dir, input_xml_dir, output_image_dir, output_xml_dir, num_augmentations=1):
    """
    对输入目录中的所有图像和对应的XML文件进行数据增强，并保存到输出目录中。
    """
    # 复制原始文件到新的输出目录
    copy_original_files(input_image_dir, input_xml_dir, output_image_dir, output_xml_dir)

    image_files = sorted([f for f in os.listdir(input_image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
    xml_files = sorted([f for f in os.listdir(input_xml_dir) if f.lower().endswith('.xml')])

    start_index = get_max_index(image_files) + 1

    for image_file, xml_file in zip(image_files, xml_files):
        image_path = os.path.join(input_image_dir, image_file)
        xml_path = os.path.join(input_xml_dir, xml_file)

        image = cv2.imread(image_path)
        h, w = image.shape[:2]

        root, objects = parse_xml(xml_path)
        bbs = [ia.BoundingBox(x1=obj['bbox'][0], y1=obj['bbox'][1], x2=obj['bbox'][2], y2=obj['bbox'][3]) for obj in
               objects]
        bbs_ia = ia.BoundingBoxesOnImage(bbs, shape=image.shape)

        # 根据num_augmentations参数控制生成增强图像的数量
        for aug_idx in range(num_augmentations):
            # 增强图像和边界框
            image_aug, bbs_aug = seq(image=image.copy(), bounding_boxes=bbs_ia)  # 使用image.copy()避免影响原图

            # 更新XML文件中的边界框信息
            updated_root = update_xml(root, objects, bbs_aug)

            # 保存增强后的图像和XML文件
            new_image_name = f"{start_index}.jpg"
            new_xml_name = f"{start_index}.xml"

            output_image_path = os.path.join(output_image_dir, new_image_name)
            output_xml_path = os.path.join(output_xml_dir, new_xml_name)

            cv2.imwrite(output_image_path, image_aug)
            save_xml(updated_root, output_xml_path)

            print(f"Processed {image_file} -> {new_image_name}")

            start_index += 1  # 每次增强后增加索引

if __name__ == "__main__":
    # 设置想要生成的增强图像数量
    num_augmentations_per_image = 2  # 每张原始图像生成5张增强图像
    augment_images_and_xmls(input_image_dir, input_xml_dir, output_image_dir, output_xml_dir,
                            num_augmentations=num_augmentations_per_image)